module diag_mod

  use spark_cubed_sphere
  use params_mod

  implicit none

  private

  public diag_type
  public diag

  type diag_panel_array_type
    ! 1: quadrature point
    ! 2: x-axis
    ! 3: y-axis
    ! 4: z-axis
    ! 5: variable
    !                            1 2 3 4 5
    real(r8), pointer, dimension(:,:,:,:,:) :: qL
    real(r8), pointer, dimension(:,:,:,:,:) :: qR
    real(r8), pointer, dimension(:,:,:,:,:) :: qT
    real(r8), pointer, dimension(:,:,:,:,:) :: qB
    real(r8), pointer, dimension(:,:,:,:,:) :: qU
    real(r8), pointer, dimension(:,:,:,:,:) :: qD
    real(r8), pointer, dimension(:,:,:,:,:) :: qQ
    real(r8), pointer, dimension(  :,:,:  ) :: u
    real(r8), pointer, dimension(  :,:,:  ) :: v
    real(r8), pointer, dimension(  :,:,:  ) :: uc
    real(r8), pointer, dimension(:,:,:,:  ) :: ucL
    real(r8), pointer, dimension(:,:,:,:  ) :: ucR
    real(r8), pointer, dimension(:,:,:,:  ) :: ucB
    real(r8), pointer, dimension(:,:,:,:  ) :: ucT
    real(r8), pointer, dimension(:,:,:,:  ) :: ucQ
    real(r8), pointer, dimension(  :,:,:  ) :: vc
    real(r8), pointer, dimension(:,:,:,:  ) :: vcL
    real(r8), pointer, dimension(:,:,:,:  ) :: vcR
    real(r8), pointer, dimension(:,:,:,:  ) :: vcB
    real(r8), pointer, dimension(:,:,:,:  ) :: vcT
    real(r8), pointer, dimension(:,:,:,:  ) :: vcQ
    real(r8), pointer, dimension(  :,:,:  ) :: w
    real(r8), pointer, dimension(  :,:,:  ) :: us
    real(r8), pointer, dimension(  :,:,:  ) :: vs
    real(r8), pointer, dimension(  :,:,:,:) :: src
    real(r8), pointer, dimension(  :,:,:,:) :: srch
    real(r8), pointer, dimension(  :,:,:,:) :: srcv
    real(r8), pointer, dimension(  :,:,:,:) :: fx
    real(r8), pointer, dimension(  :,:,:,:) :: fy
    real(r8), pointer, dimension(  :,:,:,:) :: fz
    real(r8), pointer, dimension(:,:,:,:,:) :: fxp
    real(r8), pointer, dimension(:,:,:,:,:) :: fyp
    real(r8), pointer, dimension(:,:,:,:,:) :: fzp
    real(r8), pointer, dimension(  :,:,:  ) :: vor
    logical , pointer, dimension(  :,:,:,:) :: tci
    
    ! Barotropic only
    real(r8), pointer, dimension(  :,:,:  ) :: gd ! geopotential depth
    real(r8), pointer, dimension(  :,:,:  ) :: ght! geopotential height
    real(r8), pointer, dimension(    :,:,:) :: dghdx ! Dgh / Dx
    real(r8), pointer, dimension(    :,:,:) :: dghdy ! Dgh / Dy
    
    ! Baroclinic only
    real(r8), pointer, dimension(  :,:,:  ) :: rhod   ! dry density
    real(r8), pointer, dimension(  :,:,:  ) :: rhom   ! moist density (vapor density)
    real(r8), pointer, dimension(  :,:,:  ) :: rho    ! full density
    real(r8), pointer, dimension(  :,:,:  ) :: p      ! temperature
    real(r8), pointer, dimension(  :,:,:  ) :: pt     ! potential temperature
    real(r8), pointer, dimension(  :,:,:  ) :: t      ! temperature
    real(r8), pointer, dimension(  :,:,:  ) :: mr     ! mixing ratio
    
    integer, pointer, dimension(  :,:  ) :: niter     ! iteration of implicit loop
    
  end type diag_panel_array_type

  type diag_type
    logical :: initialized = .false.
    real(r8) :: tm = 0  ! Total mass
    real(r8) :: te = 0  ! Total energy
    type(cubed_sphere_mesh_type), pointer :: mesh => null()
    type(cubed_sphere_array_type) array
    type(diag_panel_array_type) domain(6)
  contains
    procedure :: init  => diag_init
    procedure :: clear => diag_clear
    final :: diag_final
  end type diag_type

  type(diag_type) :: diag

contains

  subroutine diag_init(this, mesh)

    class(diag_type), intent(inout) :: this
    type(cubed_sphere_mesh_type), intent(in), target :: mesh

    integer ims, ime, jms, jme, kms, kme
    integer igs, ige, jgs, jge

    integer i, idom

    call this%clear()

    this%mesh => mesh
    call this%array%init(mesh)

    if (model_type == 0) then
      !!!!!!!!!!!!!!!!!
      ! Add variables !
      !!!!!!!!!!!!!!!!!
      call this%array%add_var('q1' , &
        dtype='r8', &
        long_name='Metric Jacobian * geopotential depth', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ', &
        with_halo='T:T:T:T:T:T', &
        fill_halo='T:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F')
      
      ! ------------------------------------------------------------------------
      !                                  TCI
      call this%array%add_var('tci1' , &
        dtype='l', &
        long_name='Trouble cell indicator of q1', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      ! ------------------------------------------------------------------------
      !                          Diagnostic variables
      call this%array%add_var('gd', &
        dtype='r8', &
        long_name='Geopotential depth', &
        units='m2 s-2', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('ght', &
        dtype='r8', &
        long_name='Geopotential height', &
        units='m2 s-2', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('us', &
        dtype='r8', &
        long_name='Zonal wind speed', &
        units='m s-1', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('vs', &
        dtype='r8', &
        long_name='Meridional wind speed', &
        units='m s-1', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('u', &
        dtype='r8', &
        long_name='Covariant wind speed along x-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('v', &
        dtype='r8', &
        long_name='Covariant wind speed along y-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('uc', &
        dtype='r8', &
        long_name='Contravariant wind speed along x-axis', &
        units='-', &
        loc='CA:LQ:RQ:BQ:TQ:CQ', &
        with_halo='T:T:T:T:T:T', &
        fill_halo='F:F:F:F:F:F', &
        output='F:F:F:F:F:F')

      call this%array%add_var('vc', &
        dtype='r8', &
        long_name='Contravariant wind speed along y-axis', &
        units='-', &
        loc='CA:LQ:RQ:BQ:TQ:CQ', &
        with_halo='T:T:T:T:T:T', &
        fill_halo='F:F:F:F:F:F', &
        output='F:F:F:F:F:F')

      call this%array%add_var('vor', &
        dtype='r8', &
        long_name='Relative vorticity', &
        units='s-1', &
        loc='CA', &
        with_halo='F', &
        fill_halo='F', &
        output='T')
        
      call this%array%add_var('dghdx', &
        dtype='r8', &
        long_name='Derivative of ght along x-axis', &
        units='m s-2', &
        loc='C', &
        with_halo='F', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('dghdy', &
        dtype='r8', &
        long_name='Derivative of ght along y-axis', &
        units='m s-2', &
        loc='C', &
        with_halo='F', &
        fill_halo='F', &
        output='F')

      ! ------------------------------------------------------------------------
      !                                Sources
      call this%array%add_var('src1', &
        dtype='r8', &
        long_name='Source terms of q1', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      ! ------------------------------------------------------------------------
      !                                  Flux
      call this%array%add_var('fxp1', &
        dtype='r8', &
        long_name='Flux of q1 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fyp1', &
        dtype='r8', &
        long_name='Flux of q1 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fx1', &
        dtype='r8', &
        long_name='Flux of q1 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')

      call this%array%add_var('fy1', &
        dtype='r8', &
        long_name='Flux of q1 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')

      call this%array%allocate_arrays()
      
      !!!!!!!!!!!!!!!!!
      ! Get variables !
      !!!!!!!!!!!!!!!!!
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          call this%array%get_array(this%domain(idom)%qL   , idom=idom, loc='LQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qR   , idom=idom, loc='RQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qT   , idom=idom, loc='TQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qB   , idom=idom, loc='BQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qQ   , idom=idom, loc='CQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%tci  , idom=idom, loc='C' , tag='tci')
          call this%array%get_array(this%domain(idom)%us   , idom=idom, loc='CA', var_name='us')
          call this%array%get_array(this%domain(idom)%vs   , idom=idom, loc='CA', var_name='vs')
          call this%array%get_array(this%domain(idom)%u    , idom=idom, loc='CA', var_name='u' )
          call this%array%get_array(this%domain(idom)%v    , idom=idom, loc='CA', var_name='v' )
          call this%array%get_array(this%domain(idom)%uc   , idom=idom, loc='CA', var_name='uc')
          call this%array%get_array(this%domain(idom)%ucL  , idom=idom, loc='LQ', var_name='uc')
          call this%array%get_array(this%domain(idom)%ucR  , idom=idom, loc='RQ', var_name='uc')
          call this%array%get_array(this%domain(idom)%ucB  , idom=idom, loc='BQ', var_name='uc')
          call this%array%get_array(this%domain(idom)%ucT  , idom=idom, loc='TQ', var_name='uc')
          call this%array%get_array(this%domain(idom)%ucQ  , idom=idom, loc='CQ', var_name='uc')
          call this%array%get_array(this%domain(idom)%vc   , idom=idom, loc='CA', var_name='vc')
          call this%array%get_array(this%domain(idom)%vcL  , idom=idom, loc='LQ', var_name='vc')
          call this%array%get_array(this%domain(idom)%vcR  , idom=idom, loc='RQ', var_name='vc')
          call this%array%get_array(this%domain(idom)%vcB  , idom=idom, loc='BQ', var_name='vc')
          call this%array%get_array(this%domain(idom)%vcT  , idom=idom, loc='TQ', var_name='vc')
          call this%array%get_array(this%domain(idom)%vcQ  , idom=idom, loc='CQ', var_name='vc')
          call this%array%get_array(this%domain(idom)%vor  , idom=idom, loc='CA', var_name='vor')
          call this%array%get_array(this%domain(idom)%src  , idom=idom, loc='CA', tag='src')
          call this%array%get_array(this%domain(idom)%fxp  , idom=idom, loc='LQ', tag='fxp')
          call this%array%get_array(this%domain(idom)%fyp  , idom=idom, loc='BQ', tag='fyp')
          call this%array%get_array(this%domain(idom)%fx   , idom=idom, loc='L' , tag='fx' )
          call this%array%get_array(this%domain(idom)%fy   , idom=idom, loc='B' , tag='fy' )
      
          call this%array%get_array(this%domain(idom)%gd   , idom=idom, loc='CA', var_name='gd')
          call this%array%get_array(this%domain(idom)%ght  , idom=idom, loc='CA', var_name='ght')
          call this%array%get_array(this%domain(idom)%dghdx, idom=idom, loc='C' , var_name='dghdx')
          call this%array%get_array(this%domain(idom)%dghdy, idom=idom, loc='C' , var_name='dghdy')
          this%domain(idom)%dghdx = 0
          this%domain(idom)%dghdy = 0
          
          this%domain(idom)%tci = .true.
          
          if( .not. ( trim(recon_h_scheme)=='WENO2D' .or. trim(recon_h_scheme)=='WENO2D_rhombus' ) )this%domain(idom)%tci = .false.
          
        end if
      end do
    elseif (model_type == 1) then
      !!!!!!!!!!!!!!!!!
      ! Add variables !
      !!!!!!!!!!!!!!!!!
      call this%array%add_var('q1' , &
        dtype='r8', &
        long_name='Metric Jacobian * geopotential depth', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ', &
        with_halo='T:T:T:T:T:T', &
        fill_halo='T:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F')

      call this%array%add_var('q2' , &
        dtype='r8', &
        long_name='q1 * contravariant wind component along x-axis', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ', &
        with_halo='T:T:T:T:T:T', &
        fill_halo='T:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F')

      call this%array%add_var('q3' , &
        dtype='r8', &
        long_name='q1 * contravariant wind component along y-axis', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ', &
        with_halo='T:T:T:T:T:T', &
        fill_halo='T:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F')
      
      ! ------------------------------------------------------------------------
      !                                  TCI
      call this%array%add_var('tci1' , &
        dtype='l', &
        long_name='Trouble cell indicator of q1', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      call this%array%add_var('tci2' , &
        dtype='l', &
        long_name='Trouble cell indicator of q2', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      call this%array%add_var('tci3' , &
        dtype='l', &
        long_name='Trouble cell indicator of q3', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      ! ------------------------------------------------------------------------
      !                          Diagnostic variables
      call this%array%add_var('gd', &
        dtype='r8', &
        long_name='Geopotential depth', &
        units='m2 s-2', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('ght', &
        dtype='r8', &
        long_name='Geopotential height', &
        units='m2 s-2', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('us', &
        dtype='r8', &
        long_name='Zonal wind speed', &
        units='m s-1', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('vs', &
        dtype='r8', &
        long_name='Meridional wind speed', &
        units='m s-1', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('u', &
        dtype='r8', &
        long_name='Covariant wind speed along x-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('v', &
        dtype='r8', &
        long_name='Covariant wind speed along y-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('vor', &
        dtype='r8', &
        long_name='Relative vorticity', &
        units='s-1', &
        loc='CA', &
        with_halo='F', &
        fill_halo='F', &
        output='T')
        
      call this%array%add_var('dghdx', &
        dtype='r8', &
        long_name='Derivative of ght along x-axis', &
        units='m s-2', &
        loc='C', &
        with_halo='F', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('dghdy', &
        dtype='r8', &
        long_name='Derivative of ght along y-axis', &
        units='m s-2', &
        loc='C', &
        with_halo='F', &
        fill_halo='F', &
        output='F')

      ! ------------------------------------------------------------------------
      !                                Sources
      call this%array%add_var('src1', &
        dtype='r8', &
        long_name='Source terms of q1', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      call this%array%add_var('src2', &
        dtype='r8', &
        long_name='Source terms of q2', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      call this%array%add_var('src3', &
        dtype='r8', &
        long_name='Source terms of q3', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      ! ------------------------------------------------------------------------
      !                                  Flux
      call this%array%add_var('fxp1', &
        dtype='r8', &
        long_name='Flux of q1 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fxp2', &
        dtype='r8', &
        long_name='Flux of q2 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fxp3', &
        dtype='r8', &
        long_name='Flux of q3 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fyp1', &
        dtype='r8', &
        long_name='Flux of q1 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fyp2', &
        dtype='r8', &
        long_name='Flux of q2 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fyp3', &
        dtype='r8', &
        long_name='Flux of q3 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fx1', &
        dtype='r8', &
        long_name='Flux of q1 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')

      call this%array%add_var('fx2', &
        dtype='r8', &
        long_name='Flux of q2 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')

      call this%array%add_var('fx3', &
        dtype='r8', &
        long_name='Flux of q3 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')

      call this%array%add_var('fy1', &
        dtype='r8', &
        long_name='Flux of q1 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')

      call this%array%add_var('fy2', &
        dtype='r8', &
        long_name='Flux of q2 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')

      call this%array%add_var('fy3', &
        dtype='r8', &
        long_name='Flux of q3 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')

      call this%array%allocate_arrays()
      
      !!!!!!!!!!!!!!!!!
      ! Get variables !
      !!!!!!!!!!!!!!!!!
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          call this%array%get_array(this%domain(idom)%qL   , idom=idom, loc='LQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qR   , idom=idom, loc='RQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qT   , idom=idom, loc='TQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qB   , idom=idom, loc='BQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qQ   , idom=idom, loc='CQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%tci  , idom=idom, loc='C' , tag='tci')
          call this%array%get_array(this%domain(idom)%us   , idom=idom, loc='CA', var_name='us')
          call this%array%get_array(this%domain(idom)%vs   , idom=idom, loc='CA', var_name='vs')
          call this%array%get_array(this%domain(idom)%u    , idom=idom, loc='CA', var_name='u' )
          call this%array%get_array(this%domain(idom)%v    , idom=idom, loc='CA', var_name='v' )
          call this%array%get_array(this%domain(idom)%vor  , idom=idom, loc='CA', var_name='vor')
          call this%array%get_array(this%domain(idom)%src  , idom=idom, loc='CA', tag='src')
          call this%array%get_array(this%domain(idom)%fxp  , idom=idom, loc='LQ', tag='fxp')
          call this%array%get_array(this%domain(idom)%fyp  , idom=idom, loc='BQ', tag='fyp')
          call this%array%get_array(this%domain(idom)%fx   , idom=idom, loc='L' , tag='fx' )
          call this%array%get_array(this%domain(idom)%fy   , idom=idom, loc='B' , tag='fy' )
      
          call this%array%get_array(this%domain(idom)%gd   , idom=idom, loc='CA', var_name='gd')
          call this%array%get_array(this%domain(idom)%ght  , idom=idom, loc='CA', var_name='ght')
          call this%array%get_array(this%domain(idom)%dghdx, idom=idom, loc='C' , var_name='dghdx')
          call this%array%get_array(this%domain(idom)%dghdy, idom=idom, loc='C' , var_name='dghdy')
          this%domain(idom)%dghdx = 0
          this%domain(idom)%dghdy = 0
          
          this%domain(idom)%tci = .true.
          
          if( .not. ( trim(recon_h_scheme)=='WENO2D' .or. trim(recon_h_scheme)=='WENO2D_rhombus' ) )this%domain(idom)%tci = .false.
          
        end if
      end do
    elseif(model_type==2)then
      !!!!!!!!!!!!!!!!!
      ! Add variables !
      !!!!!!!!!!!!!!!!!
      ! ------------------------------------------------------------------------
      !                            Forecast variables
      call this%array%add_var('q1' , &
        dtype='r8', &
        long_name='Metric Jacobian * dry density', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ:UQ:DQ', &
        with_halo='T:T:T:T:T:T:T:T', &
        fill_halo='F:F:F:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F:F:F')

      call this%array%add_var('q2' , &
        dtype='r8', &
        long_name='Metric Jacobian * density * contravariant wind component along x-axis', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ:UQ:DQ', &
        with_halo='T:T:T:T:T:T:T:T', &
        fill_halo='F:F:F:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F:F:F')

      call this%array%add_var('q3' , &
        dtype='r8', &
        long_name='Metric Jacobian * density * contravariant wind component along y-axis', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ:UQ:DQ', &
        with_halo='T:T:T:T:T:T:T:T', &
        fill_halo='F:F:F:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F:F:F')
      
      call this%array%add_var('q4' , &
        dtype='r8', &
        long_name='Metric Jacobian * vertical wind', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ:UQ:DQ', &
        with_halo='T:T:T:T:T:T:T:T', &
        fill_halo='F:F:F:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F:F:F:F')

      call this%array%add_var('q5' , &
        dtype='r8', &
        long_name='q1 * potential temperature', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ:UQ:DQ', &
        with_halo='T:T:T:T:T:T:T:T', &
        fill_halo='F:F:F:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F:F:F')

      call this%array%add_var('q6' , &
        dtype='r8', &
        long_name='q1 * mixing ratio', &
        units='-', &
        loc='C:CQ:LQ:RQ:TQ:BQ:UQ:DQ', &
        with_halo='T:T:T:T:T:T:T:T', &
        fill_halo='T:F:F:F:F:F:F:F', &
        tag='fcst', &
        output='F:F:F:F:F:F:F:F')
      
      ! ------------------------------------------------------------------------
      !                                  TCI
      call this%array%add_var('tci1' , &
        dtype='l', &
        long_name='Trouble cell indicator of q1', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      call this%array%add_var('tci2' , &
        dtype='l', &
        long_name='Trouble cell indicator of q2', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      call this%array%add_var('tci3' , &
        dtype='l', &
        long_name='Trouble cell indicator of q3', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')
      
      call this%array%add_var('tci4' , &
        dtype='l', &
        long_name='Trouble cell indicator of q4', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      call this%array%add_var('tci5' , &
        dtype='l', &
        long_name='Trouble cell indicator of q5', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      call this%array%add_var('tci6' , &
        dtype='l', &
        long_name='Trouble cell indicator of q6', &
        units='-', &
        loc='C', &
        with_halo='T', &
        fill_halo='F', &
        tag='tci', &
        output='F')

      ! ------------------------------------------------------------------------
      !                          Diagnostic variables
      call this%array%add_var('rhod', &
        dtype='r8', &
        long_name='Dry density', &
        units='kg m-3', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')
      
      call this%array%add_var('rhom', &
        dtype='r8', &
        long_name='Moist density', &
        units='kg m-3', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')
      
      call this%array%add_var('rho', &
        dtype='r8', &
        long_name='Full density', &
        units='kg m-3', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')
      
      call this%array%add_var('p', &
        dtype='r8', &
        long_name='Pressure', &
        units='Pa', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')
      
      call this%array%add_var('pt', &
        dtype='r8', &
        long_name='Potential Temperature', &
        units='K', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')
      
      call this%array%add_var('t', &
        dtype='r8', &
        long_name='Temperature', &
        units='K', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')
      
      call this%array%add_var('mr', &
        dtype='r8', &
        long_name='Mixing Ratio', &
        units='kg kg-1', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')
      
      call this%array%add_var('us', &
        dtype='r8', &
        long_name='Zonal wind speed', &
        units='m s-1', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('vs', &
        dtype='r8', &
        long_name='Meridional wind speed', &
        units='m s-1', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('u', &
        dtype='r8', &
        long_name='Covariant wind speed along x-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('v', &
        dtype='r8', &
        long_name='Covariant wind speed along y-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='F')

      call this%array%add_var('w', &
        dtype='r8', &
        long_name='Vertica wind speed', &
        units='m', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        output='T')

      call this%array%add_var('vor', &
        dtype='r8', &
        long_name='Relative vorticity', &
        units='s-1', &
        loc='CA', &
        with_halo='F', &
        fill_halo='F', &
        output='T')

      ! ------------------------------------------------------------------------
      !                                Sources
      call this%array%add_var('src1', &
        dtype='r8', &
        long_name='Source terms of q1', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      call this%array%add_var('src2', &
        dtype='r8', &
        long_name='Source terms of q2', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      call this%array%add_var('src3', &
        dtype='r8', &
        long_name='Source terms of q3', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      call this%array%add_var('src4', &
        dtype='r8', &
        long_name='Source terms of q4', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      call this%array%add_var('src5', &
        dtype='r8', &
        long_name='Source terms of q5', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')

      call this%array%add_var('src6', &
        dtype='r8', &
        long_name='Source terms of q6', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='src', &
        output='F')
      
      call this%array%add_var('srch1', &
        dtype='r8', &
        long_name='Source terms of q1', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srch', &
        output='F')

      call this%array%add_var('srch2', &
        dtype='r8', &
        long_name='Source terms of q2', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srch', &
        output='F')

      call this%array%add_var('srch3', &
        dtype='r8', &
        long_name='Source terms of q3', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srch', &
        output='F')

      call this%array%add_var('srch4', &
        dtype='r8', &
        long_name='Source terms of q4', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srch', &
        output='F')

      call this%array%add_var('srch5', &
        dtype='r8', &
        long_name='Source terms of q5', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srch', &
        output='F')

      call this%array%add_var('srch6', &
        dtype='r8', &
        long_name='Source terms of q6', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srch', &
        output='F')
      
      call this%array%add_var('srcv1', &
        dtype='r8', &
        long_name='Source terms of q1', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srcv', &
        output='F')

      call this%array%add_var('srcv2', &
        dtype='r8', &
        long_name='Source terms of q2', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srcv', &
        output='F')

      call this%array%add_var('srcv3', &
        dtype='r8', &
        long_name='Source terms of q3', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srcv', &
        output='F')

      call this%array%add_var('srcv4', &
        dtype='r8', &
        long_name='Source terms of q4', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srcv', &
        output='F')

      call this%array%add_var('srcv5', &
        dtype='r8', &
        long_name='Source terms of q5', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srcv', &
        output='F')

      call this%array%add_var('srcv6', &
        dtype='r8', &
        long_name='Source terms of q6', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='F', &
        tag='srcv', &
        output='F')
      
      ! ------------------------------------------------------------------------
      !                                  Flux
      call this%array%add_var('fxp1', &
        dtype='r8', &
        long_name='Flux of q1 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fxp2', &
        dtype='r8', &
        long_name='Flux of q2 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fxp3', &
        dtype='r8', &
        long_name='Flux of q3 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')
      
      call this%array%add_var('fxp4', &
        dtype='r8', &
        long_name='Flux of q4 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fxp5', &
        dtype='r8', &
        long_name='Flux of q5 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fxp6', &
        dtype='r8', &
        long_name='Flux of q6 on edge quadrature points along x-axis', &
        units='-', &
        loc='LQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fxp', &
        output='F')

      call this%array%add_var('fyp1', &
        dtype='r8', &
        long_name='Flux of q1 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fyp2', &
        dtype='r8', &
        long_name='Flux of q2 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fyp3', &
        dtype='r8', &
        long_name='Flux of q3 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fyp4', &
        dtype='r8', &
        long_name='Flux of q4 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fyp5', &
        dtype='r8', &
        long_name='Flux of q5 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fyp6', &
        dtype='r8', &
        long_name='Flux of q6 on edge quadrature points along y-axis', &
        units='-', &
        loc='BQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fyp', &
        output='F')

      call this%array%add_var('fzp1', &
        dtype='r8', &
        long_name='Flux of q1 on edge quadrature points along z-axis', &
        units='-', &
        loc='DQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fzp', &
        output='F')

      call this%array%add_var('fzp2', &
        dtype='r8', &
        long_name='Flux of q2 on edge quadrature points along z-axis', &
        units='-', &
        loc='DQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fzp', &
        output='F')

      call this%array%add_var('fzp3', &
        dtype='r8', &
        long_name='Flux of q3 on edge quadrature points along z-axis', &
        units='-', &
        loc='DQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fzp', &
        output='F')

      call this%array%add_var('fzp4', &
        dtype='r8', &
        long_name='Flux of q4 on edge quadrature points along z-axis', &
        units='-', &
        loc='DQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fzp', &
        output='F')

      call this%array%add_var('fzp5', &
        dtype='r8', &
        long_name='Flux of q5 on edge quadrature points along z-axis', &
        units='-', &
        loc='DQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fzp', &
        output='F')

      call this%array%add_var('fzp6', &
        dtype='r8', &
        long_name='Flux of q6 on edge quadrature points along z-axis', &
        units='-', &
        loc='DQ', &
        with_halo='T', &
        fill_halo='F', &
        tag='fzp', &
        output='F')
      
      call this%array%add_var('fx1', &
        dtype='r8', &
        long_name='Flux of q1 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')

      call this%array%add_var('fx2', &
        dtype='r8', &
        long_name='Flux of q2 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')

      call this%array%add_var('fx3', &
        dtype='r8', &
        long_name='Flux of q3 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')
      
      call this%array%add_var('fx4', &
        dtype='r8', &
        long_name='Flux of q4 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')

      call this%array%add_var('fx5', &
        dtype='r8', &
        long_name='Flux of q5 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')

      call this%array%add_var('fx6', &
        dtype='r8', &
        long_name='Flux of q6 along x-axis', &
        units='-', &
        loc='L', &
        with_halo='T', &
        fill_halo='F', &
        tag='fx', &
        output='F')
      
      call this%array%add_var('fy1', &
        dtype='r8', &
        long_name='Flux of q1 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')

      call this%array%add_var('fy2', &
        dtype='r8', &
        long_name='Flux of q2 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')

      call this%array%add_var('fy3', &
        dtype='r8', &
        long_name='Flux of q3 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')
      
      call this%array%add_var('fy4', &
        dtype='r8', &
        long_name='Flux of q4 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')

      call this%array%add_var('fy5', &
        dtype='r8', &
        long_name='Flux of q5 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')

      call this%array%add_var('fy6', &
        dtype='r8', &
        long_name='Flux of q6 along y-axis', &
        units='-', &
        loc='B', &
        with_halo='T', &
        fill_halo='F', &
        tag='fy', &
        output='F')
      
      call this%array%add_var('fz1', &
        dtype='r8', &
        long_name='Flux of q6 along z-axis', &
        units='-', &
        loc='D', &
        with_halo='T', &
        fill_halo='F', &
        tag='fz', &
        output='F')

      call this%array%add_var('fz2', &
        dtype='r8', &
        long_name='Flux of q6 along z-axis', &
        units='-', &
        loc='D', &
        with_halo='T', &
        fill_halo='F', &
        tag='fz', &
        output='F')

      call this%array%add_var('fz3', &
        dtype='r8', &
        long_name='Flux of q6 along z-axis', &
        units='-', &
        loc='D', &
        with_halo='T', &
        fill_halo='F', &
        tag='fz', &
        output='F')

      call this%array%add_var('fz4', &
        dtype='r8', &
        long_name='Flux of q6 along z-axis', &
        units='-', &
        loc='D', &
        with_halo='T', &
        fill_halo='F', &
        tag='fz', &
        output='F')

      call this%array%add_var('fz5', &
        dtype='r8', &
        long_name='Flux of q6 along z-axis', &
        units='-', &
        loc='D', &
        with_halo='T', &
        fill_halo='F', &
        tag='fz', &
        output='F')

      call this%array%add_var('fz6', &
        dtype='r8', &
        long_name='Flux of q6 along z-axis', &
        units='-', &
        loc='D', &
        with_halo='T', &
        fill_halo='F', &
        tag='fz', &
        output='F')

      ! ------------------------------------------------------------------------
      !                                  IMEX
      call this%array%add_var('niter', &
        dtype='i4', &
        long_name='Iteration of implicit loop', &
        units='-', &
        loc='CA', &
        with_halo='F', &
        fill_halo='F', &
        only_2d=.true., &
        output='F')
      
      call this%array%allocate_arrays()
      
      !!!!!!!!!!!!!!!!!
      ! Get variables !
      !!!!!!!!!!!!!!!!!
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          call this%array%get_array(this%domain(idom)%qL   , idom=idom, loc='LQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qR   , idom=idom, loc='RQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qT   , idom=idom, loc='TQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qB   , idom=idom, loc='BQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qU   , idom=idom, loc='UQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qD   , idom=idom, loc='DQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%qQ   , idom=idom, loc='CQ', tag='fcst')
          call this%array%get_array(this%domain(idom)%tci  , idom=idom, loc='C' , tag='tci')
          call this%array%get_array(this%domain(idom)%us   , idom=idom, loc='CA', var_name='us')
          call this%array%get_array(this%domain(idom)%vs   , idom=idom, loc='CA', var_name='vs')
          call this%array%get_array(this%domain(idom)%u    , idom=idom, loc='CA', var_name='u' )
          call this%array%get_array(this%domain(idom)%v    , idom=idom, loc='CA', var_name='v' )
          call this%array%get_array(this%domain(idom)%w    , idom=idom, loc='CA', var_name='w' )
          call this%array%get_array(this%domain(idom)%vor  , idom=idom, loc='CA', var_name='vor')
          call this%array%get_array(this%domain(idom)%src  , idom=idom, loc='CA', tag='src')
          call this%array%get_array(this%domain(idom)%srch , idom=idom, loc='CA', tag='srch')
          call this%array%get_array(this%domain(idom)%srcv , idom=idom, loc='CA', tag='srcv')
          call this%array%get_array(this%domain(idom)%fxp  , idom=idom, loc='LQ', tag='fxp')
          call this%array%get_array(this%domain(idom)%fyp  , idom=idom, loc='BQ', tag='fyp')
          call this%array%get_array(this%domain(idom)%fzp  , idom=idom, loc='DQ', tag='fzp')
          call this%array%get_array(this%domain(idom)%fx   , idom=idom, loc='L' , tag='fx' )
          call this%array%get_array(this%domain(idom)%fy   , idom=idom, loc='B' , tag='fy' )
          call this%array%get_array(this%domain(idom)%fz   , idom=idom, loc='D' , tag='fz' )
          
          call this%array%get_array(this%domain(idom)%rhod  , idom=idom, loc='CA', var_name='rhod'  )
          call this%array%get_array(this%domain(idom)%rhom  , idom=idom, loc='CA', var_name='rhom'  )
          call this%array%get_array(this%domain(idom)%rho   , idom=idom, loc='CA', var_name='rho'   )
          call this%array%get_array(this%domain(idom)%pt    , idom=idom, loc='CA', var_name='pt'    )
          call this%array%get_array(this%domain(idom)%p     , idom=idom, loc='CA', var_name='p'     )
          call this%array%get_array(this%domain(idom)%t     , idom=idom, loc='CA', var_name='t'     )
          call this%array%get_array(this%domain(idom)%mr    , idom=idom, loc='CA', var_name='mr'    )
          
          call this%array%get_array(this%domain(idom)%niter , idom=idom, loc='CA', var_name='niter' )
          
          this%domain(idom)%tci = .true.
          
          if( .not. ( trim(recon_h_scheme)=='WENO2D' .or. trim(recon_h_scheme)=='WENO2D_rhombus' ) ) this%domain(idom)%tci = .false.
          
        end if
      end do
    end if

    this%initialized = .true.

  end subroutine diag_init

  subroutine diag_clear(this)

    class(diag_type), intent(inout) :: this

    integer idom

    call this%array%clear()

    this%initialized = .false.

  end subroutine diag_clear

  subroutine diag_final(this)

    type(diag_type), intent(inout) :: this

    call this%clear()

  end subroutine diag_final

end module diag_mod
